// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/sparksql/tests/SparkFunctionBaseTest.h>

namespace kumo::pollux::functions::sparksql::test {
namespace {

class SparkPartitionIdTest : public SparkFunctionBaseTest {
 protected:
  void testSparkPartitionId(int32_t partitionId, int32_t vectorSize) {
    setSparkPartitionId(partitionId);
    auto result =
        evaluate("spark_partition_id()", make_row_vector(ROW({}), vectorSize));
    ASSERT_TRUE(result->is_constant_encoding());
    pollux::test::assertEqualVectors(
        make_constant(partitionId, vectorSize), result);
  }
};

TEST_F(SparkPartitionIdTest, basic) {
  testSparkPartitionId(0, 1);
  testSparkPartitionId(100, 1);
  testSparkPartitionId(0, 100);
  testSparkPartitionId(100, 100);
}

TEST_F(SparkPartitionIdTest, error) {
  auto row_vector = make_row_vector(ROW({}), 1);

  queryCtx_->testingOverrideConfigUnsafe({{}});
  POLLUX_ASSERT_THROW(
      evaluate("spark_partition_id()", row_vector),
      "Spark partition id is not set");

  setSparkPartitionId(-1);
  POLLUX_ASSERT_THROW(
      evaluate("spark_partition_id()", row_vector),
      "Invalid Spark partition id");
}
} // namespace
} // namespace kumo::pollux::functions::sparksql::test
